package momo.multitree.algo;

import momo.multitree.structure.*;

import java.util.LinkedHashMap;
import java.util.LinkedHashSet;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;


public class PrimMinSpanTree {
	
	private Log log = LogFactory.getLog(PrimMinSpanTree.class);
	
	/**
	 * assumes costs are symmetric
	 * 
	 * @param g
	 * @param root
	 * @return
	 */
	public Tree mstCostTree(Graph g, Node root) {
		assert g.isSymmetricCost();
		
		Tree mst = new Tree(root);
		LinkedHashSet<Node> visited = new LinkedHashSet<Node>();
		visited.add(root);
		
		LinkedHashSet<Edge> neigh = new LinkedHashSet<Edge>();
		LinkedHashMap<Node, Double> distances = new LinkedHashMap<Node, Double>();
		
		for (Edge edge : g.getEdgesOfNode(root)) {
			neigh.add(edge);
			distances.put(edge.nextNode(root), edge.getCost(root));
		}
		
		while (visited.size() < g.getNoOfNodes()) {
			double minCost = Double.POSITIVE_INFINITY;
			double minDist = Double.POSITIVE_INFINITY;
			Edge minEdge = null;
			
			for (Edge edge : neigh) {
				Node outsideNode = edge.getHeadNode();
				if (visited.contains(outsideNode))
					outsideNode = edge.getTailNode();
				if (!visited.contains(outsideNode)) {
					double currCost = edge.getCost(edge.getHeadNode());
					if ( currCost < minCost) {
						minCost = currCost;
						minEdge = edge;
						minDist = distances.get(outsideNode);
					} else {
						if ( currCost == minCost ) {
							double currDist = distances.get(outsideNode);
							if (currDist < minDist) {
								minCost = currCost;
								minEdge = edge;
								minDist = currDist;
							}
						}
					}
				}
			}
			log.debug("minEdge: "+minEdge.toString());
			
			Node headNode = minEdge.getHeadNode();
			Node tailNode = minEdge.getTailNode();
			Node newNode = null;
			if (visited.contains(headNode))
				newNode = tailNode;
			else
				newNode = headNode;
			log.debug("new node: "+newNode.toString());
			
			visited.add(newNode);
			mst.add(newNode);
			mst.add(minEdge);
			
			for (Edge edge : g.getEdgesOfNode(newNode)) {
				neigh.add(edge);
				Node other = edge.nextNode(newNode);
				if (!visited.contains(other)) {
					if (!distances.containsKey(other))
						distances.put(other, distances.get(newNode) + edge.getCost(newNode));
					else {
						double prevDist = distances.get(other);
						double currDist = distances.get(newNode) + edge.getCost(newNode);
						if (currDist < prevDist)
							distances.put(other, currDist);								
					}
				}
			}
		}
		
		return mst;
		
	}
	
}
